from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from datetime import datetime
from config import DATABASE_URL

class DatabaseManager:
    def __init__(self):
        self.engine: Engine = create_engine(DATABASE_URL)
    
    def _convert_value(self, value):
        """
        转换特殊类型的值
        """
        if isinstance(value, datetime):
            return value.isoformat()
        return value
    
    def execute_query(self, query: str) -> list:
        """
        执行SQL查询并返回结果
        """
        with self.engine.connect() as connection:
            result = connection.execute(text(query))
            # 获取列名
            columns = result.keys()
            # 将结果转换为字典列表，同时处理特殊类型
            return [
                {
                    col: self._convert_value(value)
                    for col, value in zip(columns, row)
                }
                for row in result
            ]
    
    def get_tables(self) -> list:
        """
        获取所有表名及其注释
        """
        query = """
        SELECT 
            t.table_name,
            obj_description((quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass, 'pg_class') as table_comment
        FROM information_schema.tables t
        WHERE t.table_schema = 'public' 
        ORDER BY t.table_name;
        """
        return self.execute_query(query)
    
    def get_table_columns(self, table_name: str) -> list:
        """
        获取指定表的所有列信息及其注释
        """
        query = f"""
        SELECT 
            c.column_name,
            c.data_type,
            c.is_nullable,
            col_description((quote_ident(c.table_schema) || '.' || quote_ident(c.table_name))::regclass::oid, c.ordinal_position) as column_comment
        FROM information_schema.columns c
        WHERE c.table_name = '{table_name}'
        ORDER BY c.ordinal_position;
        """
        return self.execute_query(query)
    
    def get_table_data(self, table_name: str, limit: int = 10) -> list:
        """
        获取指定表的数据
        """
        query = f"""
        SELECT *
        FROM {table_name}
        LIMIT {limit};
        """
        return self.execute_query(query) 